/**
* @file allocator.cpp
*
* Copyright (C) Huawei Technologies Co., Ltd. 2023-2023. All Rights Reserved.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
*/

#include <memory>

#include "log_inner.h"
#include "acl/acl_rt_allocator.h"
#include "model/acl_resource_manager.h"
#include "toolchain/profiling_manager.h"
#include "toolchain/resource_statistics.h"
#include "framework/memory/allocator_desc.h"

aclrtAllocatorDesc aclrtAllocatorCreateDesc()
{
    ACL_PROFILING_REG(acl::AclProfType::AclrtAllocatorCreateDesc);
    ACL_ADD_APPLY_TOTAL_COUNT(acl::ACL_STATISTICS_CREATE_DESTROY_ALLOCATOR_DESC);
    ACL_LOG_INFO("Create allocator description.");
    ge::AllocatorDesc *allocatorDesc = new(std::nothrow) ge::AllocatorDesc;
    if (allocatorDesc == nullptr) {
        ACL_LOG_INNER_ERROR("alloc ge::AllocatorDesc memory failed");
        return nullptr;
    }
    ACL_ADD_APPLY_SUCCESS_COUNT(acl::ACL_STATISTICS_CREATE_DESTROY_ALLOCATOR_DESC);
    return static_cast<void *>(allocatorDesc);
}

aclError aclrtAllocatorDestroyDesc(aclrtAllocatorDesc allocatorDesc)
{
    ACL_LOG_INFO("Destroy allocator description, allocatorDesc %p.", allocatorDesc);
    ACL_PROFILING_REG(acl::AclProfType::AclrtAllocatorDestroyDesc);
    ACL_ADD_RELEASE_TOTAL_COUNT(acl::ACL_STATISTICS_CREATE_DESTROY_ALLOCATOR_DESC);
    ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocatorDesc);
    delete static_cast<ge::AllocatorDesc *>(allocatorDesc);
    allocatorDesc = nullptr;
    ACL_ADD_RELEASE_SUCCESS_COUNT(acl::ACL_STATISTICS_CREATE_DESTROY_ALLOCATOR_DESC);
    return ACL_SUCCESS;
}

aclError aclrtAllocatorSetObjToDesc(aclrtAllocatorDesc allocatorDesc, aclrtAllocator allocator)
{
    ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocatorDesc);
    ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocator);
    ACL_LOG_INFO("Set allocator to allocator description, allocatorDesc %p.", allocatorDesc);
    static_cast<ge::AllocatorDesc *>(allocatorDesc)->obj = allocator;
    return ACL_SUCCESS;
}

aclError aclrtAllocatorSetAllocFuncToDesc(aclrtAllocatorDesc allocatorDesc, aclrtAllocatorAllocFunc func)
{
    ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocatorDesc);
    ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(func);
    ACL_LOG_INFO("Set alloc function to allocator description, allocatorDesc %p.", allocatorDesc);
    static_cast<ge::AllocatorDesc *>(allocatorDesc)->alloc_func = func;
    return ACL_SUCCESS;
}

aclError aclrtAllocatorSetFreeFuncToDesc(aclrtAllocatorDesc allocatorDesc, aclrtAllocatorFreeFunc func)
{
    ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocatorDesc);
    ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(func);
    ACL_LOG_INFO("Set free function to allocator description, allocatorDesc %p.", allocatorDesc);
    static_cast<ge::AllocatorDesc *>(allocatorDesc)->free_func = func;
    return ACL_SUCCESS;
}

aclError aclrtAllocatorSetAllocAdviseFuncToDesc(aclrtAllocatorDesc allocatorDesc, aclrtAllocatorAllocAdviseFunc func)
{
    ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocatorDesc);
    ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(func);
    ACL_LOG_INFO("Set alloc advise function to allocator description, allocatorDesc %p.", allocatorDesc);
    static_cast<ge::AllocatorDesc *>(allocatorDesc)->alloc_advise_func = func;
    return ACL_SUCCESS;
}

aclError aclrtAllocatorSetGetAddrFromBlockFuncToDesc(aclrtAllocatorDesc allocatorDesc,
                                                     aclrtAllocatorGetAddrFromBlockFunc func)
{
    ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocatorDesc);
    ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(func);
    ACL_LOG_INFO("Set get_addr_from_block function to allocator description, allocatorDesc %p.", allocatorDesc);
    static_cast<ge::AllocatorDesc *>(allocatorDesc)->get_addr_from_block_func = func;
    return ACL_SUCCESS;
}

aclError aclrtAllocatorRegister(aclrtStream stream, aclrtAllocatorDesc allocatorDesc)
{
    // stream must be not null when register external allocator
    ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(stream);
    ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(allocatorDesc);

    const auto ret = acl::AclResourceManager::GetInstance().CreateExternalAllocator(stream, allocatorDesc);
    if (ret != ACL_SUCCESS) {
        ACL_LOG_INNER_ERROR("Register external allocator failed, stream %p, allocatorDesc %p.", stream, allocatorDesc);
        return ACL_ERROR_INVALID_PARAM;
    }
    ACL_LOG_INFO("Register external allocator success, stream %p, allocatorDesc %p.", stream, allocatorDesc);
    return ACL_SUCCESS;
}

aclError aclrtAllocatorUnregister(aclrtStream stream)
{
    ACL_REQUIRES_NOT_NULL_WITH_INPUT_REPORT(stream);

    ACL_LOG_INFO("Unregister external allocator success, stream %p.", stream);
    return acl::AclResourceManager::GetInstance().DeleteExternalAllocator(stream);
}
